package com.shujia.mllib

import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.sql.expressions.UserDefinedFunction
import org.apache.spark.sql.{DataFrame, SaveMode, SparkSession}

object Demo3ImageMakeData {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()

    import spark.implicits._
    import org.apache.spark.sql.functions._


    //1、读取图片的数据
    val imageDF: DataFrame = spark.read
      .format("image")
      .load("E:\\data\\train")

    imageDF.printSchema()


    //编写自定义函数处理图片数据
    val comData: UserDefinedFunction = udf((data: Array[Byte]) => {
      val comData: Array[Double] = data.map(b => {
        //将图片中的像素点转换成0和1，白色:1,黑色：0
        if (b.toInt < 0) {
          1.0
        } else {
          0.0
        }
      })
      //转换成向量返回
      Vectors.dense(comData).toSparse
    })

    //获取图片名称
    val getName: UserDefinedFunction = udf((origin: String) => {
      origin.split("/").last
    })

    val nameAndFeatures: DataFrame = imageDF
      //取出文件名和文件的数据
      .select($"image.origin" as "origin", $"image.data" as "data")
      //解析图片的名称和数据
      .select(getName($"origin") as "name", comData($"data") as "features")

    //读取图片的名称和图片中的数字
    val nameAndLabel: DataFrame = spark.read
      .format("csv")
      .option("sep", " ")
      .schema("name STRING, label DOUBLE")
      .load("spark/data/image_res.txt")


    //关联获取目标值和特征向量
    val joinDF: DataFrame = nameAndFeatures.join(nameAndLabel, "name")

    val tranDF: DataFrame = joinDF.select($"label", $"features")

    //保存处理的结果
    tranDF
      .write
      .format("libsvm")
      .mode(SaveMode.Overwrite)
      .save("spark/data/images")

  }

}
